# models.py
import torch

class PhysicsInformedNN(torch.nn.Module):
    '''PINNs模型'''

    def __init__(self):
        super().__init__()
        self.layers = torch.nn.Sequential(
            torch.nn.Linear(1, 128),
            torch.nn.ReLU(),
            torch.nn.Linear(128, 128),
            torch.nn.ReLU(),
            torch.nn.Linear(128, 64),
            torch.nn.ReLU(),
            torch.nn.Linear(64, 1)
        )
        self.apply(self._init_weights)

    def _init_weights(self, m: torch.nn.Module):
        if isinstance(m, torch.nn.Linear):
            torch.nn.init.kaiming_normal_(m.weight, nonlinearity="relu")
            torch.nn.init.zeros_(m.bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.layers(x)

    @staticmethod
    def physics_equation(F_pred_raw: torch.Tensor, A: torch.Tensor, 
                         E: torch.Tensor, g31: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        return (g31 * F_pred_raw * t) / (A * E)